import numpy as np
import librosa

def calculate_snr(original_audio, watermarked_audio):
    if len(original_audio) != len(watermarked_audio):
        raise ValueError("两个音频信号长度必须相同")
    signal_power = np.sum(original_audio ** 2)
    noise_power = np.sum((original_audio - watermarked_audio) ** 2)
    if noise_power == 0:
        return float('inf')
    return 10 * np.log10(signal_power / noise_power)

# 读取音频文件
original_audio, _ = librosa.load('original.wav', sr=None)
watermarked_audio, _ = librosa.load('output.wav', sr=None)

try:
    snr = calculate_snr(original_audio, watermarked_audio)
    print(f"信噪比 (SNR): {snr} dB")
except ValueError as e:
    print(f"错误: {e}")